import torch
import cv2
import numpy as np
import tkinter as tk
from tkinter import filedialog, messagebox, ttk
from PIL import Image, ImageTk
from model import CustomNet
from torchvision import transforms
import os
import time
from tkinter import font as tkfont


class HandGestureRecognizer:
    def __init__(self, model_path='./models/best_model.pth', num_classes=6):
        self.model = CustomNet(num_classes=num_classes)

        try:
            pretrained_dict = torch.load(model_path, map_location=torch.device('cpu'))
            model_dict = self.model.state_dict()
            pretrained_dict = {k: v for k, v in pretrained_dict.items()
                              if k in model_dict and v.shape == model_dict[k].shape}
            model_dict.update(pretrained_dict)
            self.model.load_state_dict(model_dict)
            print(f"成功加载模型权重，匹配参数: {len(pretrained_dict)}/{len(model_dict)}")
        except Exception as e:
            print(f"模型加载失败: {e}")
            print("使用随机初始化权重（仅调试用，实际应确保 best_model.pth 存在）")

        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        self.model.to(self.device)
        self.model.eval()

        self.transform = transforms.Compose([
            transforms.Resize((224, 224)),
            transforms.ToTensor(),
            transforms.Normalize(mean=[0.485, 0.456, 0.406],
                                 std=[0.229, 0.224, 0.225])
        ])

        self.root = tk.Tk()
        self.root.title("手势数字识别系统 (完整适配版)")
        self.root.geometry("900x650")
        self.root.configure(bg="#1a1a2e")
        self.root.resizable(True, True)

        self.title_font = tkfont.Font(family="Microsoft YaHei UI", size=24, weight="bold")
        self.subtitle_font = tkfont.Font(family="Microsoft YaHei UI", size=16, weight="bold")
        self.result_font = tkfont.Font(family="Microsoft YaHei UI", size=20, weight="bold")
        self.normal_font = tkfont.Font(family="Microsoft YaHei UI", size=12)

        self._create_advanced_ui()

    def _create_advanced_ui(self):
        title_frame = tk.Frame(self.root, bg="#16213e", height=80)
        title_frame.pack(fill=tk.X)
        title_frame.pack_propagate(False)

        title_label = tk.Label(title_frame, text="手势数字识别系统",
                               font=self.title_font, fg="#e94560", bg="#16213e")
        title_label.place(relx=0.5, rely=0.5, anchor=tk.CENTER)

        main_frame = tk.Frame(self.root, bg="#1a1a2e")
        main_frame.pack(fill=tk.BOTH, expand=True, padx=20, pady=10)

        self.image_frame = tk.Frame(main_frame, bg="#2b2d42",
                                    bd=3, relief=tk.GROOVE, highlightthickness=0)
        self.image_frame.pack(side=tk.LEFT, fill=tk.BOTH, expand=True, padx=5, pady=5)

        self.image_label = tk.Label(self.image_frame, text="\n\n请上传手势图片\n进行识别\n\n",
                                    font=self.subtitle_font, fg="#a8dadc", bg="#2b2d42",
                                    justify=tk.CENTER)
        self.image_label.pack(fill=tk.BOTH, expand=True, padx=20, pady=20)

        upload_btn = tk.Button(self.image_frame, text="上传图片",
                               font=self.normal_font, fg="#f1faee", bg="#457b9d",
                               activeforeground="#f1faee", activebackground="#1d3557",
                               bd=0, padx=15, pady=8,
                               command=self._upload_image)
        upload_btn.pack(pady=15)
        upload_btn.bind("<Enter>", lambda e: upload_btn.config(bg="#3a86ff"))
        upload_btn.bind("<Leave>", lambda e: upload_btn.config(bg="#457b9d"))

        result_frame = tk.Frame(main_frame, bg="#1a1a2e", width=350)
        result_frame.pack(side=tk.RIGHT, fill=tk.Y, padx=5, pady=5)
        result_frame.pack_propagate(False)

        result_title = tk.Label(result_frame, text="识别结果",
                                font=self.subtitle_font, fg="#e94560", bg="#1a1a2e")
        result_title.pack(pady=(20, 10))

        self.result_display = tk.Frame(result_frame, bg="#2b2d42",
                                       bd=2, relief=tk.SUNKEN)
        self.result_display.pack(fill=tk.BOTH, expand=True, padx=20, pady=10)

        self.finger_count = tk.Label(self.result_display, text="--",
                                     font=self.result_font, fg="#4cc9f0", bg="#2b2d42")
        self.finger_count.pack(pady=20)

        self.confidence_label = tk.Label(self.result_display, text="置信度: 0%",
                                         font=self.normal_font, fg="#a8dadc", bg="#2b2d42")
        self.confidence_label.pack(pady=5)

        history_frame = tk.Frame(result_frame, bg="#1a1a2e")
        history_frame.pack(fill=tk.X, padx=20, pady=(20, 10))

        history_title = tk.Label(history_frame, text="识别历史",
                                 font=self.normal_font, fg="#4cc9f0", bg="#1a1a2e")
        history_title.pack(anchor=tk.W)

        self.history_listbox = tk.Listbox(history_frame,
                                          font=self.normal_font,
                                          bg="#2b2d42", fg="#f1faee",
                                          width=30, height=5,
                                          selectbackground="#457b9d")
        self.history_listbox.pack(side=tk.LEFT, fill=tk.BOTH, expand=True)

        scrollbar = tk.Scrollbar(history_frame, command=self.history_listbox.yview)
        scrollbar.pack(side=tk.RIGHT, fill=tk.Y)
        self.history_listbox.config(yscrollcommand=scrollbar.set)

        status_frame = tk.Frame(self.root, bg="#16213e", height=30)
        status_frame.pack(side=tk.BOTTOM, fill=tk.X)
        status_frame.pack_propagate(False)

        self.status_label = tk.Label(status_frame, text="就绪",
                                     font=self.normal_font, fg="#f1faee", bg="#16213e")
        self.status_label.place(relx=0.5, rely=0.5, anchor=tk.CENTER)

    def _upload_image(self):
        file_path = filedialog.askopenfilename(
            title="选择手势图片",
            filetypes=[("图像文件", "*.png;*.jpg;*.jpeg;*.bmp")]
        )
        if not file_path:
            return

        self.status_label.config(text=f"正在处理: {os.path.basename(file_path)}")
        self.root.update()

        try:
            pil_image = Image.open(file_path).convert('RGB')
            image_np = np.array(pil_image)
            input_tensor = self.transform(pil_image).unsqueeze(0).to(self.device)
        except Exception as e:
            messagebox.showerror("图像加载失败", f"无法读取图片: {str(e)}")
            self.status_label.config(text="识别失败")
            return

        try:
            with torch.no_grad():
                outputs = self.model(input_tensor)
                probs = torch.nn.functional.softmax(outputs, dim=1)
                confidence, predicted = torch.max(probs, 1)
        except Exception as e:
            messagebox.showerror("模型推理失败", f"推理时发生错误: {str(e)}")
            self.status_label.config(text="识别失败")
            return

        self._display_image(image_np)
        self._animate_result(predicted.item(), confidence.item() * 100)

        timestamp = time.strftime("%H:%M:%S", time.localtime())
        # 关键修改点：调整历史记录格式，去掉手指数量信息
        history_entry = f"[{timestamp}] (置信度: {confidence.item() * 100:.1f}%)"
        self.history_listbox.insert(0, history_entry)
        if self.history_listbox.size() > 10:
            self.history_listbox.delete(10)

        self.status_label.config(text="识别完成")

    def _display_image(self, image):
        h, w = image.shape[:2]
        max_display_size = min(self.image_frame.winfo_width() - 40,
                               self.image_frame.winfo_height() - 80)
        if max_display_size > 0:
            scale = max_display_size / max(h, w)
            new_size = (int(w * scale), int(h * scale))
            pil_img = Image.fromarray(image).resize(new_size, Image.Resampling.LANCZOS)
            image = np.array(pil_img)

        pil_img_tk = ImageTk.PhotoImage(Image.fromarray(image))
        self.image_label.config(image=pil_img_tk, text="", fg="#a8dadc")
        self.image_label.image = pil_img_tk

    def _animate_result(self, finger_count, conf_percent):
        def update_count(count):
            if count <= finger_count:
                self.finger_count.config(text=str(count), fg="#4cc9f0")
                self.root.after(50, update_count, count + 1)
            else:
                def update_conf(conf):
                    if conf <= conf_percent:
                        self.confidence_label.config(text=f"置信度: {conf:.1f}%")
                        self.root.after(20, update_conf, conf + 1)
                    else:
                        self.finger_count.config(fg="#e94560")
                update_conf(0)
        update_count(0)

    def run(self):
        self.root.mainloop()


if __name__ == "__main__":
    app = HandGestureRecognizer(num_classes=6)
    app.run()